import numpy as np
import matplotlib.pyplot as plt
from channel import awgn_torch
import commpy
from commpy.channelcoding.ldpc import get_ldpc_code_params, ldpc_bp_decode, triang_ldpc_systematic_encode
from channel_codec.ldpc_codec import LDPC_v2
from communication import *
import torch
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

ldpc_codec = LDPC_v2(
    ldpc_design_file=r"F:\research-code\2025\H256+LDPC\channel_codec\1440.720.txt"
)


def commpy_version(bits):
    ldpc_encode_bits = ldpc_codec.encode(bits)
    # ldpc_bits_with_noise = ldpc_encode_bits
    msg_bits = ldpc_encode_bits
    # ldpc_bits_with_noise = commpy.awgn(ldpc_encode_bits.reshape(-1), e)
    msg_bits = 1 - 2 * msg_bits
    snr_linear = 10 ** (e / 10)
    sigma2 = 1 / (2 * 1 * snr_linear)
    sigma = torch.sqrt(torch.tensor(sigma2))
    noise = (sigma * torch.randn_like(torch.tensor(msg_bits), dtype=torch.float32)).numpy()
    ldpc_bits_with_noise = msg_bits + noise
    ps = np.mean(np.abs(msg_bits)**2)
    pn = np.mean(np.abs(noise)**2)
    print(f"Commpy Ps: {ps}, Pn:{pn}, SNR: {10*np.log10(ps/pn)}")
    ldpc_bits_with_noise = 2 * ldpc_bits_with_noise / sigma2
    ldpc_decode_bits = ldpc_codec.decode(ldpc_bits_with_noise)
    # bers_commpy.append((bits != ldpc_decode_bits).sum() / len(bits))
    print(f"Commpy version SNR: {e} dB, BER: {(bits != ldpc_decode_bits).sum() / len(bits)}")
    return (bits != ldpc_decode_bits).sum() / len(bits)


def torch_version(bits):
    bits = torch.tensor(bits, device=device, dtype=torch.int)
    ldpc_encode_bits = triang_ldpc_systematic_encode_torch(bits, ldpc_params)
    msg_bits = ldpc_encode_bits
    msg_bits = 1 - 2 * msg_bits
    snr_linear = 10 ** (e / 10)
    sigma2 = 1 / (2 * 1 * snr_linear)
    sigma = torch.sqrt(torch.tensor(sigma2))
    noise = (sigma * torch.randn_like(msg_bits, dtype=torch.float32))
    ldpc_bits_with_noise = msg_bits + noise
    ps = torch.mean(torch.abs(msg_bits)**2, dtype=torch.float32)
    pn = torch.mean(torch.abs(noise)**2, dtype=torch.float32)
    print(f"Torch Ps: {ps}, Pn:{pn}, SNR: {10*torch.log10(ps/pn)}")
    ldpc_bits_with_noise = 2 * ldpc_bits_with_noise / sigma2
    ldpc_bits_with_noise = torch.tensor(ldpc_bits_with_noise.detach().cpu().numpy().reshape(-1, order='F').astype(float),
                                    device=device)
    ldpc_decode_bits = torch_ldpc_bp_decode(ldpc_bits_with_noise, ldpc_params, "MSA", 10)[0][:720]
    # bers_torch.append((bits.detach().cpu().numpy() != ldpc_decode_bits.detach().cpu().numpy().reshape(-1, order='F')[:len(bits)]).sum() / len(bits))
    print(f"Torch version SNR: {e} dB, BER: {(bits.detach().cpu().numpy() != ldpc_decode_bits.detach().cpu().numpy().reshape(-1, order='F')[:len(bits)]).sum() / len(bits)}")
    return (bits.detach().cpu().numpy() != ldpc_decode_bits.detach().cpu().numpy().reshape(-1, order='F')[:len(bits)]).sum() / len(bits)


device = "cuda"
bits = np.random.randint(0, 2, 1440*100)
ldpc_params = get_ldpc_code_params(ldpc_design_filename=r"F:\research-code\2025\H256+LDPC\channel_codec\1440.720.txt", compute_matrix=True)

# 示例
ebno_dbs = np.arange(-5, 26, 1)
bers_commpy, bers_torch = [], []

for e in ebno_dbs:
    print(f"SNR: {e} dB")
    # TODO:Commpy version
    ber_com = commpy_version(bits)

    # TODO: Torch version
    ber_tor = torch_version(bits)
    bers_commpy.append(ber_com)
    bers_torch.append(ber_tor)

plt.figure(figsize=(18, 6))
plt.subplot(1, 2, 1)
plt.plot(ebno_dbs, bers_commpy, "ro-")
plt.title("1/2LDPC Commpy version")
plt.xlabel("SNR(dB)")
plt.ylabel("BER")
plt.subplot(1, 2, 2)
plt.plot(ebno_dbs, bers_torch, "bo-")
plt.title("1/2LDPC Torch version")
plt.xlabel("SNR(dB)")
plt.ylabel("BER")
plt.show()
